import uuid
import streamlit as st
import io
import re
from pydub import AudioSegment
from biz.integrations.aliyun.speech_recognition_api import recognize_audio

from biz.config.vanna_config import (
    generate_questions_cached,
    generate_sql_cached,
    run_sql_cached,
    generate_plotly_code_cached,
    generate_plot_cached,
    generate_followup_cached,
    should_generate_chart_cached,
    is_sql_valid_cached,
    generate_summary_cached
)

avatar_url = "static/telchina.svg"

st.set_page_config(page_title="Telchina AI for SQL",
                   initial_sidebar_state="collapsed",
                   page_icon=":shark:",
                   layout="wide")

# 使用 markdown 插入自定义 CSS 来隐藏右上角的菜单
hide_menu_style = """
    <style>
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    </style>
"""
st.markdown(hide_menu_style, unsafe_allow_html=True)

# 自定义CSS
st.markdown(
    """
    <style>
    /* 设置侧边栏的宽度 */
    .eczjsme18 {
        width: 245px !important;
    }
    
    .e12wn80j15 {
        position: fixed;
        bottom: 5px;
        width: 10%;
        background-color: rgba(255, 255, 255, 0) ;
        padding: 0px;
        display: flex;
        justify-content: space-between;
        z-index: 9999;
        height: 50px;
        margin-left: -15px;
    }
    
    .e12wn80j14 {
        background-color: rgba(255, 255, 255, 0);
        height: 50px;
    }
    
    .e2wxzia0 {
        display: none;
    }
    
    .stAppDeployButton {
        display: none;
    }
    
    .e1nzilvr3 {
        margin-top: -90px
    }
    
    </style>
    """,
    unsafe_allow_html=True
)

st.sidebar.title("输出设置")
st.sidebar.checkbox("显示SQL", value=True, key="show_sql")
st.sidebar.checkbox("显示表格", value=True, key="show_table")
#st.sidebar.checkbox("Show Plotly Code", value=True, key="show_plotly_code")
st.sidebar.checkbox("显示图表", value=True, key="show_chart")
#st.sidebar.checkbox("Show Summary", value=True, key="show_summary")
#st.sidebar.checkbox("Show Follow-up Questions", value=True, key="show_followup")
st.sidebar.button("Reset", on_click=lambda: set_question(None), use_container_width=True)

st.title("Telchina AI for SQL")
# st.sidebar.write(st.session_state)

def remove_angle_brackets_content(text):
    # 使用正则表达式匹配尖括号及其包含的内容
    pattern = r'<[^>]+>'
    # 替换匹配到的内容为空字符串
    cleaned_text = re.sub(pattern, '', text)
    return cleaned_text

def set_question(question):
    st.session_state["my_question"] = question
    st.session_state["messages"] = []

def sql_run(sql,my_question):
    try :
        if sql:
            if is_sql_valid_cached(sql=sql):
                if st.session_state.get("show_sql", True):
                    assistant_message_sql = st.chat_message(
                        "assistant", avatar=avatar_url
                    )
                    assistant_message_sql.code(sql, language="sql", line_numbers=True)
            else:
                assistant_message = st.chat_message(
                    "assistant", avatar=avatar_url
                )
                assistant_message.write(sql)
                st.stop()

            df = run_sql_cached(sql=sql)

            if df is not None:
                st.session_state["df"] = df

            if st.session_state.get("df") is not None:
                if st.session_state.get("show_table", True):
                    df = st.session_state.get("df")
                    assistant_message_table = st.chat_message(
                        "assistant",
                        avatar=avatar_url,
                    )
                    if len(df) > 20:
                        assistant_message_table.text("前20行数据")
                        assistant_message_table.dataframe(df.head(20))
                    else:
                        assistant_message_table.dataframe(df)

                if should_generate_chart_cached(question=my_question, sql=sql, df=df):
                    code = generate_plotly_code_cached(question=my_question, sql=sql, df=df)

                    if st.session_state.get("show_plotly_code", False):
                        assistant_message_plotly_code = st.chat_message(
                            "assistant",
                            avatar=avatar_url,
                        )
                        assistant_message_plotly_code.code(
                            code, language="python", line_numbers=True
                        )

                    if code is not None and code != "":
                        if st.session_state.get("show_chart", True):
                            assistant_message_chart = st.chat_message(
                                "assistant",
                                avatar=avatar_url,
                            )
                            fig = generate_plot_cached(code=code, df=df)
                            if fig is not None:
                                assistant_message_chart.plotly_chart(fig, key=uuid.uuid4())
                            else:
                                assistant_message_chart.error("I couldn't generate a chart")

                if st.session_state.get("show_summary", False):
                    assistant_message_summary = st.chat_message(
                        "assistant",
                        avatar=avatar_url,
                    )
                    summary = generate_summary_cached(question=my_question, df=df)
                    if summary is not None:
                        assistant_message_summary.text(summary)

                if st.session_state.get("show_followup", False):
                    assistant_message_followup = st.chat_message(
                        "assistant",
                        avatar=avatar_url,
                    )
                    followup_questions = generate_followup_cached(
                        question=my_question, sql=sql, df=df
                    )
                    st.session_state["df"] = None

                    if len(followup_questions) > 0:
                        assistant_message_followup.text(
                            "Here are some possible follow-up questions"
                        )
                        # Print the first 5 follow-up questions
                        for question in followup_questions[:5]:
                            assistant_message_followup.button(question, on_click=set_question, args=(question,))

        else:
            assistant_message_error = st.chat_message(
                "assistant", avatar=avatar_url
            )
            assistant_message_error.error("I wasn't able to generate SQL for that question")
    except Exception as e:
        print(f"发生了异常：{e}")

if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "assistant", "type": "text", "content": "嗨，我是一个可以搜索SQL的机器人。我能帮你什么?"}
    ]

for msg in st.session_state.messages:
    if msg["type"] == "sql" and is_sql_valid_cached(sql=msg["content"]):
        #st.chat_message(msg["role"], avatar=avatar_url).write(msg["content"])
        sql_run(msg["content"], msg["question"])
    elif msg["type"] == "text":
        if msg["role"] == "user":
            st.chat_message(msg["role"]).write(msg["content"])
        else:
            st.chat_message(msg["role"], avatar=avatar_url).write(msg["content"])
    else:
        st.chat_message(msg["role"], avatar=avatar_url).write(msg["content"])

if 'audio_value_changed' not in st.session_state:
    st.session_state.audio_value_changed = False
# 定义回调函数
def audio_callback():
    st.session_state.audio_value_changed = True


my_question = st.chat_input("问我一个关于你的数据的问题",)

audio_value = st.experimental_audio_input("Record a voice message", on_change=audio_callback, label_visibility="collapsed")
if audio_value and st.session_state.audio_value_changed:
    try:
        st.session_state.audio_value_changed = False
        # 创建一个BytesIO对象来模拟文件
        audio_buffer = io.BytesIO(audio_value.getvalue())
        audio = AudioSegment.from_file(audio_buffer, format="wav")
        # 更改采样率为 16000 Hz
        audio = audio.set_frame_rate(16000)
        # 将修改后的音频保存为新的字节数据
        with io.BytesIO() as output:
            audio.export(output, format="wav")
            audio_bytes = output.getvalue()
        text = recognize_audio(audio_bytes, num_processes=1, timeout=30)
        print("识别结果:", text)
        my_question = remove_angle_brackets_content(text)
    except Exception as e:
        print(f"发生了异常：{e}")



if my_question:
    st.session_state.messages.append({"role": "user", "type":"text", "content": my_question})
    st.session_state["my_question"] = my_question
    user_message = st.chat_message("user")
    user_message.write(f"{my_question}")

    question_sql_history = []
    question_sql_history = [
        (item["question"], item["content"]) 
        for item in st.session_state.messages
        if item["type"] == "sql"
    ]
    
    # Only store the last MAX_QUESTION_HISTORY questions
    if len(question_sql_history) > 10:
        question_sql_history = question_sql_history[-10:]

    sql = generate_sql_cached(question=my_question, question_sql_history=question_sql_history)
    st.session_state.messages.append({"role": "assistant", "type":"sql", "question":my_question, "content": sql})
    sql_run(sql,my_question)





